import torch.nn as nn

from mmdet.models import HEADS, SSDHead


def SeperableConv2d(in_channels,
                    out_channels,
                    kernel_size=1,
                    stride=1,
                    padding=0):
    """Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d.
    """
    return nn.Sequential(
        nn.Conv2d(in_channels=in_channels,
                  out_channels=in_channels,
                  kernel_size=kernel_size,
                  groups=in_channels,
                  stride=stride,
                  padding=padding),
        nn.BatchNorm2d(in_channels),
        nn.ReLU6(),
        nn.Conv2d(in_channels=in_channels,
                  out_channels=out_channels,
                  kernel_size=1),
    )


@HEADS.register_module()
class SSDLiteHead(SSDHead):
    def __init__(self, **kwargs):
        super(SSDLiteHead, self).__init__(**kwargs)
        num_anchors = self.anchor_generator.num_base_anchors
        reg_convs = list()
        cls_convs = list()
        for i, channel in enumerate(kwargs["in_channels"]):
            reg_convs.append(
                SeperableConv2d(in_channels=channel,
                                out_channels=num_anchors[i] * 4,
                                kernel_size=3,
                                stride=1,
                                padding=1))
            cls_convs.append(
                SeperableConv2d(in_channels=channel,
                                out_channels=num_anchors[i] *
                                (kwargs["num_classes"] + 1),
                                kernel_size=3,
                                stride=1,
                                padding=1))
        self.reg_convs = nn.ModuleList(reg_convs)
        self.cls_convs = nn.ModuleList(cls_convs)
